from openai import OpenAI
import os

# 从环境变量中获取API密钥
api_key = os.getenv("SPARK_API_KEY")


def send_message(client, model, messages):
    """
    发送消息并获取流式响应。

    Args:
        client (OpenAI): OpenAI客户端。
        model (str): 使用的模型名称。
        messages (list): 对话消息列表。

    Returns:
        str: 助手的完整回复。
    """
    # 发送请求并获取流式响应
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        stream=True,
    )
    response_content = ""
    # 处理流式响应
    for chunk in response:
        if chunk.choices and chunk.choices[0].delta and chunk.choices[0].delta.content:
            response_content += chunk.choices[0].delta.content
            print(chunk.choices[0].delta.content, end='', flush=True)
    return response_content


def main():
    """
    主函数。
    """
    # 设置API基础URL和模型名称
    base_url = 'https://spark-api-open.xf-yun.com/v1'
    model = 'lite'

    # 初始化OpenAI客户端
    client = OpenAI(
        api_key=api_key,
        base_url=base_url,
    )

    # 初始对话历史
    messages = [
        {"role": "system", "content": "你是脱口秀演员，以专业的角度讲笑梗。"},
    ]

    while True:
        # 发送消息并获取响应
        response_content = send_message(client, model, messages)

        # 将响应添加到对话历史
        messages.append({"role": "assistant", "content": response_content})

        # 获取用户输入
        user_input = input("\n你: ")
        if user_input.lower() in ["exit", "quit"]:
            break

        # 将用户输入添加到对话历史
        messages.append({"role": "user", "content": user_input})


if __name__ == "__main__":
    main()
